"""
初始化RBAC角色和权限数据
运行方式：python -m tool.init_rbac_data
"""
import asyncio
from models.m import Role, Permission, RolePermissionRelation, User, UserRoleRelation


async def init_rbac():
    """初始化RBAC基础数据"""
    
    print("=" * 50)
    print("开始初始化RBAC权限管理系统...")
    print("=" * 50)
    
    # ========== 1. 创建角色 ==========
    print("\n1. 创建角色...")
    roles_data = [
        {"id": 1, "name": "普通用户", "code": "user", "description": "普通用户角色"},
        {"id": 2, "name": "作者", "code": "author", "description": "小说作者角色"},
        {"id": 3, "name": "管理员", "code": "admin", "description": "系统管理员角色"},
    ]
    
    for role_data in roles_data:
        role_data_copy = role_data.copy()  # 创建副本，避免修改原字典
        code = role_data_copy.pop("code")  # 取出 code 作为查询条件
        defaults = role_data_copy  # 剩余数据作为创建时的默认值
        role, created = await Role.get_or_create(
            code=code,
            defaults=defaults
        )
        if created:
            print(f"  [OK] 创建角色: {role.name} ({role.code})")
        else:
            print(f"  [-] 角色已存在: {role.name} ({role.code})")
    
    # ========== 2. 创建权限 ==========
    print("\n2. 创建权限...")
    permissions_data = [
        # 小说相关权限
        {"id": 1, "name": "创建小说", "code": "novel:create", "module": "novel", "description": "创建新小说的权限"},
        {"id": 2, "name": "更新小说", "code": "novel:update", "module": "novel", "description": "更新小说信息的权限"},
        {"id": 3, "name": "删除小说", "code": "novel:delete", "module": "novel", "description": "删除小说的权限"},
        {"id": 4, "name": "发布小说", "code": "novel:publish", "module": "novel", "description": "发布小说的权限"},
        {"id": 5, "name": "审核小说", "code": "novel:audit", "module": "novel", "description": "审核小说的权限"},
        # 章节相关权限
        {"id": 6, "name": "创建章节", "code": "chapter:create", "module": "chapter", "description": "创建章节的权限"},
        {"id": 7, "name": "更新章节", "code": "chapter:update", "module": "chapter", "description": "更新章节的权限"},
        {"id": 8, "name": "删除章节", "code": "chapter:delete", "module": "chapter", "description": "删除章节的权限"},
        {"id": 9, "name": "发布章节", "code": "chapter:publish", "module": "chapter", "description": "发布章节的权限"},
        {"id": 10, "name": "审核章节", "code": "chapter:audit", "module": "chapter", "description": "审核章节的权限"},
        # 管理相关权限
        {"id": 11, "name": "用户管理", "code": "user:manage", "module": "user", "description": "管理用户信息的权限"},
        {"id": 12, "name": "角色管理", "code": "role:manage", "module": "role", "description": "管理角色的权限"},
        {"id": 13, "name": "权限管理", "code": "permission:manage", "module": "permission", "description": "管理权限的权限"},
    ]
    
    for perm_data in permissions_data:
        perm_data_copy = perm_data.copy()  # 创建副本，避免修改原字典
        code = perm_data_copy.pop("code")  # 取出 code 作为查询条件
        defaults = perm_data_copy  # 剩余数据作为创建时的默认值
        perm, created = await Permission.get_or_create(
            code=code,
            defaults=defaults
        )
        if created:
            print(f"  [OK] 创建权限: {perm.name} ({perm.code})")
        else:
            print(f"  [-] 权限已存在: {perm.name} ({perm.code})")
    
    # ========== 3. 为角色分配权限 ==========
    print("\n3. 为角色分配权限...")
    
    # 获取数据库连接（使用原生SQL避免主键问题）
    from tortoise import connections
    db = connections.get("default")
    
    # 获取角色
    try:
        author_role = await Role.get(code="author")
        admin_role = await Role.get(code="admin")
    except Exception as e:
        print(f"  [ERROR] 获取角色失败: {e}")
        return
    
    # 作者角色的权限
    author_permissions = [
        "novel:create", "novel:update", "novel:delete", "novel:publish",
        "chapter:create", "chapter:update", "chapter:delete", "chapter:publish",
    ]
    
    print(f"\n  为 {author_role.name} 角色分配权限:")
    
    for perm_code in author_permissions:
        try:
            perm = await Permission.get(code=perm_code)
            # 检查是否已存在
            existing = await RolePermissionRelation.filter(
                role_id=author_role.id,
                permission_id=perm.id
            ).first()
            if not existing:
                # 使用原生 SQL 插入，不包含 id 字段，让数据库自动生成
                await db.execute_query(
                    f"INSERT INTO role_permission_relations (role_id, permission_id, created_at) "
                    f"VALUES ({author_role.id}, {perm.id}, NOW())"
                )
                print(f"    [OK] {perm.name}")
            else:
                print(f"    [-] {perm.name} (已存在)")
        except Exception as e:
            print(f"    [ERROR] 权限 {perm_code} 出错: {str(e)}")
    
    # 管理员角色的权限（所有权限）
    admin_permissions = [
        "novel:create", "novel:update", "novel:delete", "novel:publish", "novel:audit",
        "chapter:create", "chapter:update", "chapter:delete", "chapter:publish", "chapter:audit",
        "user:manage", "role:manage", "permission:manage",
    ]
    
    print(f"\n  为 {admin_role.name} 角色分配权限:")
    for perm_code in admin_permissions:
        try:
            perm = await Permission.get(code=perm_code)
            # 检查是否已存在
            existing = await RolePermissionRelation.filter(
                role_id=admin_role.id,
                permission_id=perm.id
            ).first()
            if not existing:
                # 使用原生 SQL 插入，不包含 id 字段，让数据库自动生成
                await db.execute_query(
                    f"INSERT INTO role_permission_relations (role_id, permission_id, created_at) "
                    f"VALUES ({admin_role.id}, {perm.id}, NOW())"
                )
                print(f"    [OK] {perm.name}")
            else:
                print(f"    [-] {perm.name} (已存在)")
        except Exception as e:
            print(f"    [ERROR] 权限 {perm_code} 出错: {str(e)}")
    
    # ========== 4. 根据现有 users 表的 role 字段，为用户分配角色 ==========
    print("\n4. 根据现有用户角色分配RBAC角色...")
    
    # 获取角色对象
    user_role = await Role.get(code="user")
    author_role_obj = await Role.get(code="author")
    admin_role_obj = await Role.get(code="admin")
    
    # 导入 UserRole 枚举
    from models.m import UserRole
    
    # 为 role=1 (USER) 的用户分配 'user' 角色
    users_with_role_1 = await User.filter(role=UserRole.USER).all()
    count = 0
    for user in users_with_role_1:
        existing = await UserRoleRelation.filter(
            user_id=user.id,
            role_id=user_role.id
        ).first()
        if not existing:
            # 使用原生 SQL 插入，不包含 id 字段，让数据库自动生成
            await db.execute_query(
                f"INSERT INTO user_role_relations (user_id, role_id, created_at) "
                f"VALUES ({user.id}, {user_role.id}, NOW())"
            )
            count += 1
    if count > 0:
        print(f"  [OK] 为 {count} 个普通用户分配了 'user' 角色")
    
    # 为 role=2 (AUTHOR) 的用户分配 'author' 角色
    users_with_role_2 = await User.filter(role=UserRole.AUTHOR).all()
    count = 0
    for user in users_with_role_2:
        existing = await UserRoleRelation.filter(
            user_id=user.id,
            role_id=author_role_obj.id
        ).first()
        if not existing:
            # 使用原生 SQL 插入，不包含 id 字段，让数据库自动生成
            await db.execute_query(
                f"INSERT INTO user_role_relations (user_id, role_id, created_at) "
                f"VALUES ({user.id}, {author_role_obj.id}, NOW())"
            )
            count += 1
    if count > 0:
        print(f"  [OK] 为 {count} 个作者分配了 'author' 角色")
    
    # 为 role=3 (ADMIN) 的用户分配 'admin' 角色
    users_with_role_3 = await User.filter(role=UserRole.ADMIN).all()
    count = 0
    for user in users_with_role_3:
        existing = await UserRoleRelation.filter(
            user_id=user.id,
            role_id=admin_role_obj.id
        ).first()
        if not existing:
            # 使用原生 SQL 插入，不包含 id 字段，让数据库自动生成
            await db.execute_query(
                f"INSERT INTO user_role_relations (user_id, role_id, created_at) "
                f"VALUES ({user.id}, {admin_role_obj.id}, NOW())"
            )
            count += 1
    if count > 0:
        print(f"  [OK] 为 {count} 个管理员分配了 'admin' 角色")
    
    print("\n" + "=" * 50)
    print("RBAC权限管理系统初始化完成！")
    print("=" * 50)


async def main():
    """主函数"""
    from tortoise import Tortoise
    from tool.sessing import TORTOISE_ORM
    from tortoise import connections
    
    # 初始化数据库连接
    await Tortoise.init(config=TORTOISE_ORM)
    await Tortoise.generate_schemas()
    
    # 修复表结构：确保 id 字段是 AUTO_INCREMENT
    db = connections.get("default")
    try:
        # 修复 role_permission_relations 表
        await db.execute_query(
            "ALTER TABLE role_permission_relations MODIFY COLUMN id BIGINT NOT NULL AUTO_INCREMENT"
        )
        print("修复 role_permission_relations 表结构成功")
    except Exception as e:
        if "doesn't exist" not in str(e) and "Unknown column" not in str(e):
            print(f"修复 role_permission_relations 表结构: {e}")
    
    try:
        # 修复 user_role_relations 表
        await db.execute_query(
            "ALTER TABLE user_role_relations MODIFY COLUMN id BIGINT NOT NULL AUTO_INCREMENT"
        )
        print("修复 user_role_relations 表结构成功")
    except Exception as e:
        if "doesn't exist" not in str(e) and "Unknown column" not in str(e):
            print(f"修复 user_role_relations 表结构: {e}")
    
    try:
        await init_rbac()
    finally:
        await Tortoise.close_connections()


if __name__ == "__main__":
    asyncio.run(main())

